﻿// Copyright (c) 2018 Siegfried Pammer
// 
// Permission is hereby granted, free of charge, to any person obtaining a copy of this
// software and associated documentation files (the "Software"), to deal in the Software
// without restriction, including without limitation the rights to use, copy, modify, merge,
// publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons
// to whom the Software is furnished to do so, subject to the following conditions:
// 
// The above copyright notice and this permission notice shall be included in all copies or
// substantial portions of the Software.
// 
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
// INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
// PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE
// FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
// OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.

using System;
using System.Collections.Generic;
using System.Collections.Immutable;
using System.Diagnostics;
using System.IO;
using System.IO.Compression;
using System.Linq;
using System.Reflection.Metadata;
using System.Reflection.Metadata.Ecma335;
using System.Reflection.PortableExecutable;
using System.Security.Cryptography;
using System.Text;
using ICSharpCode.Decompiler.CSharp;
using ICSharpCode.Decompiler.CSharp.OutputVisitor;
using ICSharpCode.Decompiler.CSharp.Syntax;
using ICSharpCode.Decompiler.CSharp.TypeSystem;
using ICSharpCode.Decompiler.IL;
using ICSharpCode.Decompiler.Metadata;
using ICSharpCode.Decompiler.TypeSystem;
using ICSharpCode.Decompiler.Util;

namespace ICSharpCode.Decompiler.DebugInfo
{
	public class PortablePdbWriter
	{
		static readonly Guid CSharpLanguageGuid = new Guid("3f5162f8-07c6-11d3-9053-00c04fa302a1");

		static readonly Guid DebugInfoEmbeddedSource = new Guid("0e8a571b-6926-466e-b4ad-8ab04611f5fe");
		static readonly Guid MethodSteppingInformationBlobId = new Guid("54FD2AC5-E925-401A-9C2A-F94F171072F8");

		static readonly Guid HashAlgorithmSHA1 = new Guid("ff1816ec-aa5e-4d10-87f7-6f4963833460");
		static readonly Guid HashAlgorithmSHA256 = new Guid("8829d00f-11b8-4213-878b-770e8597ac16");
		static readonly FileVersionInfo decompilerVersion = FileVersionInfo.GetVersionInfo(typeof(CSharpDecompiler).Assembly.Location);

		public static bool HasCodeViewDebugDirectoryEntry(PEFile file)
		{
			return file.Reader.ReadDebugDirectory().Any(entry => entry.Type == DebugDirectoryEntryType.CodeView);
		}

		public static void WritePdb(PEFile file, CSharpDecompiler decompiler, DecompilerSettings settings, Stream targetStream, bool noLogo = false)
		{
			MetadataBuilder metadata = new MetadataBuilder();
			MetadataReader reader = file.Metadata;
			var entrypointHandle = MetadataTokens.MethodDefinitionHandle(file.Reader.PEHeaders.CorHeader.EntryPointTokenOrRelativeVirtualAddress);

			var sequencePointBlobs = new Dictionary<MethodDefinitionHandle, (DocumentHandle Document, BlobHandle SequencePoints)>();
			var emptyList = new List<SequencePoint>();
			var stateMachineMethods = new List<(MethodDefinitionHandle MoveNextMethod, MethodDefinitionHandle KickoffMethod)>();
			var customDocumentDebugInfo = new List<(DocumentHandle Parent, GuidHandle Guid, BlobHandle Blob)>();
			var customMethodDebugInfo = new List<(MethodDefinitionHandle Parent, GuidHandle Guid, BlobHandle Blob)>();
			var globalImportScope = metadata.AddImportScope(default, default);

			foreach (var handle in reader.GetTopLevelTypeDefinitions()) {
				var type = reader.GetTypeDefinition(handle);

				// Generate syntax tree
				var syntaxTree = decompiler.DecompileTypes(new[] { handle });
				if (!syntaxTree.HasChildren)
					continue;

				// Generate source and checksum
				if (!noLogo)
					syntaxTree.InsertChildAfter(null, new Comment(" PDB and source generated by ICSharpCode.Decompiler " + decompilerVersion.FileVersion), Roles.Comment);
				var sourceText = SyntaxTreeToString(syntaxTree, settings);

				// Generate sequence points for the syntax tree
				var sequencePoints = decompiler.CreateSequencePoints(syntaxTree);

				// Generate other debug information
				var debugInfoGen = new DebugInfoGenerator(decompiler.TypeSystem);
				syntaxTree.AcceptVisitor(debugInfoGen);

				lock (metadata) {
					var sourceBlob = WriteSourceToBlob(metadata, sourceText, out var sourceCheckSum);
					var name = metadata.GetOrAddDocumentName(type.GetFullTypeName(reader).ReflectionName.Replace('.', Path.DirectorySeparatorChar) + ".cs");

					// Create Document(Handle)
					var document = metadata.AddDocument(name,
						hashAlgorithm: metadata.GetOrAddGuid(HashAlgorithmSHA256),
						hash: metadata.GetOrAddBlob(sourceCheckSum),
						language: metadata.GetOrAddGuid(CSharpLanguageGuid));

					// Add embedded source to the PDB
					customDocumentDebugInfo.Add((document,
						metadata.GetOrAddGuid(DebugInfoEmbeddedSource),
						sourceBlob));

					debugInfoGen.Generate(metadata, globalImportScope);

					foreach (var function in debugInfoGen.Functions) {
						var method = function.MoveNextMethod ?? function.Method;
						var methodHandle = (MethodDefinitionHandle)method.MetadataToken;
						sequencePoints.TryGetValue(function, out var points);
						ProcessMethod(methodHandle, document, points, syntaxTree);
						if (function.MoveNextMethod != null) {
							stateMachineMethods.Add((
								(MethodDefinitionHandle)function.MoveNextMethod.MetadataToken,
								(MethodDefinitionHandle)function.Method.MetadataToken
							));
						}
						if (function.IsAsync) {
							customMethodDebugInfo.Add((methodHandle,
								metadata.GetOrAddGuid(MethodSteppingInformationBlobId),
								metadata.GetOrAddBlob(function.AsyncDebugInfo.BuildBlob(methodHandle))));
						}
					}
				}
			}

			foreach (var method in reader.MethodDefinitions) {
				var md = reader.GetMethodDefinition(method);

				if (sequencePointBlobs.TryGetValue(method, out var info)) {
					metadata.AddMethodDebugInformation(info.Document, info.SequencePoints);
				} else {
					metadata.AddMethodDebugInformation(default, default);
				}
			}

			stateMachineMethods.SortBy(row => MetadataTokens.GetRowNumber(row.MoveNextMethod));
			foreach (var row in stateMachineMethods) {
				metadata.AddStateMachineMethod(row.MoveNextMethod, row.KickoffMethod);
			}
			customMethodDebugInfo.SortBy(row => MetadataTokens.GetRowNumber(row.Parent));
			foreach (var row in customMethodDebugInfo) {
				metadata.AddCustomDebugInformation(row.Parent, row.Guid, row.Blob);
			}
			customDocumentDebugInfo.SortBy(row => MetadataTokens.GetRowNumber(row.Parent));
			foreach (var row in customDocumentDebugInfo) {
				metadata.AddCustomDebugInformation(row.Parent, row.Guid, row.Blob);
			}

			var debugDir = file.Reader.ReadDebugDirectory().FirstOrDefault(dir => dir.Type == DebugDirectoryEntryType.CodeView);
			var portable = file.Reader.ReadCodeViewDebugDirectoryData(debugDir);
			var contentId = new BlobContentId(portable.Guid, debugDir.Stamp);
			PortablePdbBuilder serializer = new PortablePdbBuilder(metadata, GetRowCounts(reader), entrypointHandle, blobs => contentId);
			BlobBuilder blobBuilder = new BlobBuilder();
			serializer.Serialize(blobBuilder);
			blobBuilder.WriteContentTo(targetStream);

			void ProcessMethod(MethodDefinitionHandle method, DocumentHandle document,
				List<SequencePoint> sequencePoints, SyntaxTree syntaxTree)
			{
				var methodDef = reader.GetMethodDefinition(method);
				int localSignatureRowId;
				MethodBodyBlock methodBody;
				if (methodDef.RelativeVirtualAddress != 0) {
					methodBody = file.Reader.GetMethodBody(methodDef.RelativeVirtualAddress);
					localSignatureRowId = methodBody.LocalSignature.IsNil ? 0 : MetadataTokens.GetRowNumber(methodBody.LocalSignature);
				} else {
					methodBody = null;
					localSignatureRowId = 0;
				}
				if (sequencePoints?.Count > 0)
					sequencePointBlobs.Add(method, (document, EncodeSequencePoints(metadata, localSignatureRowId, sequencePoints)));
				else
					sequencePointBlobs.Add(method, (default, default));
			}
		}

		static BlobHandle WriteSourceToBlob(MetadataBuilder metadata, string sourceText, out byte[] sourceCheckSum)
		{
			var builder = new BlobBuilder();
			using (var memory = new MemoryStream()) {
				var deflate = new DeflateStream(memory, CompressionLevel.Optimal, leaveOpen: true);
				byte[] bytes = Encoding.UTF8.GetBytes(sourceText);
				deflate.Write(bytes, 0, bytes.Length);
				deflate.Close();
				byte[] buffer = memory.ToArray();
				builder.WriteInt32(bytes.Length); // compressed
				builder.WriteBytes(buffer);
				using (var hasher = SHA256.Create()) {
					sourceCheckSum = hasher.ComputeHash(bytes);
				}
			}

			return metadata.GetOrAddBlob(builder);
		}

		static BlobHandle EncodeSequencePoints(MetadataBuilder metadata, int localSignatureRowId, List<SequencePoint> sequencePoints)
		{
			if (sequencePoints.Count == 0)
				return default;
			var writer = new BlobBuilder();
			// header:
			writer.WriteCompressedInteger(localSignatureRowId);

			int previousOffset = -1;
			int previousStartLine = -1;
			int previousStartColumn = -1;

			for (int i = 0; i < sequencePoints.Count; i++) {
				var sequencePoint = sequencePoints[i];
				// delta IL offset:
				if (i > 0)
					writer.WriteCompressedInteger(sequencePoint.Offset - previousOffset);
				else
					writer.WriteCompressedInteger(sequencePoint.Offset);
				previousOffset = sequencePoint.Offset;

				if (sequencePoint.IsHidden) {
					writer.WriteInt16(0);
					continue;
				}

				int lineDelta = sequencePoint.EndLine - sequencePoint.StartLine;
				int columnDelta = sequencePoint.EndColumn - sequencePoint.StartColumn;

				writer.WriteCompressedInteger(lineDelta);

				if (lineDelta == 0) {
					writer.WriteCompressedInteger(columnDelta);
				} else {
					writer.WriteCompressedSignedInteger(columnDelta);
				}

				if (previousStartLine < 0) {
					writer.WriteCompressedInteger(sequencePoint.StartLine);
					writer.WriteCompressedInteger(sequencePoint.StartColumn);
				} else {
					writer.WriteCompressedSignedInteger(sequencePoint.StartLine - previousStartLine);
					writer.WriteCompressedSignedInteger(sequencePoint.StartColumn - previousStartColumn);
				}

				previousStartLine = sequencePoint.StartLine;
				previousStartColumn = sequencePoint.StartColumn;
			}

			return metadata.GetOrAddBlob(writer);
		}

		static ImmutableArray<int> GetRowCounts(MetadataReader reader)
		{
			var builder = ImmutableArray.CreateBuilder<int>(MetadataTokens.TableCount);
			for (int i = 0; i < MetadataTokens.TableCount; i++) {
				builder.Add(reader.GetTableRowCount((TableIndex)i));
			}

			return builder.MoveToImmutable();
		}

		static string SyntaxTreeToString(SyntaxTree syntaxTree, DecompilerSettings settings)
		{
			StringWriter w = new StringWriter();
			TokenWriter tokenWriter = new TextWriterTokenWriter(w);
			syntaxTree.AcceptVisitor(new InsertParenthesesVisitor { InsertParenthesesForReadability = true });
			tokenWriter = TokenWriter.WrapInWriterThatSetsLocationsInAST(tokenWriter);
			syntaxTree.AcceptVisitor(new CSharpOutputVisitor(tokenWriter, settings.CSharpFormattingOptions));
			return w.ToString();
		}
	}
}
